/*
 * Copyright 2014 Semidata. All rights reserved.
 */
package com.semidata.trp.decision.utils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @author shewei@semidata.com, She Wei
 *
 */
public class SampleUtils {
	
	private static class Node<T> {
		public double weight;
		public T value;
		public double totalWeight;
		
		public Node(double weight, T value, double totalWeight) {
			this.weight = weight;
			this.value = value;
			this.totalWeight = totalWeight;
		}
	}
	
	
	private static <T> List<Node<T>> rwsHeap(Map<T, Integer> items) {
		List<Node<T>> result = new ArrayList<>();
		result.add(null);
		
		for (Map.Entry<T, Integer> entry : items.entrySet()) {
			result.add(new Node<T>(entry.getValue(), entry.getKey(), entry.getValue()));
		}
		
		for (int i = result.size() - 1; i > 1; i--) {
			result.get(i >> 1).totalWeight += result.get(i).totalWeight;
		}
		return result;
	}
		
	private static <T> T rwsHeapPop(List<Node<T>> items) {
		
		double gas = items.get(1).totalWeight * Math.random();
		
		int i = 1;
		
		while (gas > items.get(i).weight) {
			gas -= items.get(i).weight;
			i <<= 1;
			if (gas > items.get(i).totalWeight) {
				gas -= items.get(i).totalWeight;
				i ++;
			}
		}
		
		double weight = items.get(i).weight;
		T value = items.get(i).value;
		items.get(i).weight = 0;
		while (i > 0) {
			items.get(i).totalWeight -= weight;
			i >>= 1;
		}
		return value;
	}
	
	public static <T> List<T> weightedSample(Map<T, Integer> items, int n) {
		if (items == null || items.isEmpty() || n <= 0 || items.size() < n)
			return null;
		
		List<Node<T>> heap = rwsHeap(items);
		
		List<T> result = new ArrayList<T>();
		for (int i = 0; i < n; i++) {
			result.add(rwsHeapPop(heap));
		}
		return result;
	}

}
